package cn.cas.mango.util.keywordsBatch;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

class UnionFind {
  private Map<String, String> parent;

  public UnionFind(Set<String> ids) {
    parent = new HashMap<>();
    for (String id : ids) {
      parent.put(id, id);
    }
  }

  public String find(String x) {
    if (!parent.get(x).equals(x)) {
      parent.put(x, find(parent.get(x)));
    }
    return parent.get(x);
  }

  public void union(String x, String y) {
    String rootX = find(x);
    String rootY = find(y);
    if (!rootX.equals(rootY)) {
      parent.put(rootX, rootY);
    }
  }

  public Map<String, List<String>> getClusters() {
    Map<String, List<String>> clusters = new HashMap<>();
    for (String id : parent.keySet()) {
      String root = find(id);
      clusters.computeIfAbsent(root, k -> new ArrayList<>()).add(id);
    }
    return clusters;
  }
}
